#include <cub/cub.cuh>

__global__ void radixSortKernel(int* data, int* output, int numItems) {
    // 使用 cub::KeyValuePair 进行排序
    typedef cub::KeyValuePair<int, int> KeyValuePair;

    // 分配存储空间
    cub::CachingDeviceAllocator allocator;
    KeyValuePair* temp_storage = nullptr;
    size_t temp_storage_bytes = 0;
    cub::DeviceRadixSort::SortKeys(temp_storage, temp_storage_bytes, data, output, numItems);

    cudaMalloc(&temp_storage, temp_storage_bytes);

    // 执行基数排序
    cub::DeviceRadixSort::SortKeys(temp_storage, temp_storage_bytes, data, output, numItems);

    // 释放临时存储空间
    cudaFree(temp_storage);
}

int main() {
    // 设置输入数据
    const int numItems = 100;
    int* d_data, d_out_put;
    cudaMalloc(&d_data, numItems * sizeof(int));
    cudaMalloc(&d_out_put, numItems * sizeof(int));

    // 将数据传输到 GPU
    int h_data[numItems] = { /* initialize your data */ };
    cudaMemcpy(d_data, h_data, numItems * sizeof(int), cudaMemcpyHostToDevice);

    // 调用核函数
    radixSortKernel<<<1, 1>>>(d_data, d_output, numItems);

    // 将排序后的数据传输回主机
    cudaMemcpy(h_data, d_output, numItems * sizeof(int), cudaMemcpyDeviceToHost);

    // 在这里处理排序后的数据

    // 释放 GPU 上的内存
    cudaFree(d_data);

    return 0;
}
